#include "Platform/OpenGL/openGLShader.h"
// #include "Core/time.h"

#include <fstream>
// #include <glad/glad.h>

#include <glm/gtc/type_ptr.hpp>

#include <iostream>
#include "Core/utils.h"
#include "Core/timeCore.h"

#include <shaderc/shaderc.hpp>

namespace CC
{
	namespace Utils
	{
		static const std::string GetCacheDirectory()
		{
			// TODO: make sure the assets directory is valid
			return PATHCONFIG("assets/cache/shader/opengl");
		}

		static const std::string GetDefaultDirectory()
		{
			// TODO: make sure the assets directory is valid
			return PATHCONFIG("assets/shaders");
		}

		// static void CreateCacheDirectoryIfNeeded()
		// {
		// 	std::string cacheDirectory = GetCacheDirectory();
		// 	if (!std::filesystem::exists(cacheDirectory))
		// 		std::filesystem::create_directories(cacheDirectory);
		// }

		static const char* GLShaderStageCachedOpenGLFileExtension(uint32_t stage)
		{
			switch (stage)
			{
				case GL_VERTEX_SHADER:    return ".cached_opengl.vert";
				case GL_FRAGMENT_SHADER:  return ".cached_opengl.frag";
			}
			CC_ASSERT(false, "ASSERT: GLShaderStageCachedOpenGLFileExtension\n");
			return "";
		}

		static const char* GLShaderStageOpenGLFileExtension(uint32_t stage)
		{
			switch (stage)
			{
				case GL_VERTEX_SHADER:    return ".vert";
				case GL_FRAGMENT_SHADER:  return ".frag";
			}
			CC_ASSERT(false, "ASSERT: GLShaderStageOpenGLFileExtension\n");
			return "";
		}

		static shaderc_shader_kind GLShaderStageToShaderC(GLenum stage)
		{
			switch (stage)
			{
				case GL_VERTEX_SHADER:   return shaderc_glsl_vertex_shader;
				case GL_FRAGMENT_SHADER: return shaderc_glsl_fragment_shader;
			}
			CC_ASSERT(false, "ASSERT: GLShaderStageToShaderC\n");
			return (shaderc_shader_kind)0;
		}

		static const char* GLShaderStageToString(GLenum stage)
		{
			switch (stage)
			{
				case GL_VERTEX_SHADER:   return "GL_VERTEX_SHADER";
				case GL_FRAGMENT_SHADER: return "GL_FRAGMENT_SHADER";
			}
			CC_ASSERT(false, "ASSERT: GLShaderStageToString\n");
			return nullptr;
		}

	}

	OpenGLShader::OpenGLShader(const std::string &name, const std::string &vertexFile, const std::string &fragmentFile)
		: m_Name(name)
	{
		std::unordered_map<GLenum, std::string> sources;
		sources[GL_VERTEX_SHADER] = vertexFile;
		sources[GL_FRAGMENT_SHADER] = fragmentFile;

		Timer timer;
        const int HasPreCompile = 0;
        if(HasPreCompile)
		{
			CompileOrGetOpenGLBinaries(sources);
			CreateProgram();
		}
        else{
			 CompileOnly(Utils::GetDefaultDirectory() + "/" + vertexFile + Utils::GLShaderStageOpenGLFileExtension(GL_VERTEX_SHADER), 
			 		Utils::GetDefaultDirectory() + "/" + fragmentFile + Utils::GLShaderStageOpenGLFileExtension(GL_FRAGMENT_SHADER));
		}
		
		CC_CORE_WARN("Shader creation took %f ms\n", timer.ElapsedMillis());
		CC_CORE_INFO("%s Shader creation, id = %d\n", name.c_str(), m_RendererID);
	}

	OpenGLShader::~OpenGLShader()
	{

		glDeleteProgram(m_RendererID);
	}

	void OpenGLShader::CompileOrGetOpenGLBinaries(std::unordered_map<GLenum, std::string> &shaderSources)
	{
		shaderc::Compiler compiler;
		shaderc::CompileOptions options;

		options.SetTargetEnvironment(shaderc_target_env_opengl, shaderc_env_version_opengl_4_5);
		const bool optimize = false;
		if (optimize)
			options.SetOptimizationLevel(shaderc_optimization_level_performance);

		std::string cacheDirectory = Utils::GetCacheDirectory();

		auto &shaderData = m_OpenGLSPIRV;

		shaderData.clear();
		m_OpenGLSourceCode.clear();

	    std::unordered_map<GLenum, std::string>::iterator it = shaderSources.begin();
		for(it; it != shaderSources.end(); it++)
		{
			GLenum stage = it->first;
			std::string source = it->second;

			std::string cachedPath = cacheDirectory + "/" + source + Utils::GLShaderStageCachedOpenGLFileExtension(stage);
			CC_CORE_INFO("cachedPath = %s\n", cachedPath.c_str());

			std::ifstream in(cachedPath, std::ios::in | std::ios::binary);
			if (in.is_open())
			{
				in.seekg(0, std::ios::end);
				auto size = in.tellg();
				in.seekg(0, std::ios::beg);

				auto &data = shaderData[stage];
				data.resize(size / sizeof(uint32_t));
				in.read((char *)data.data(), size);
			}
			else
			{
				m_FilePath = Utils::GetDefaultDirectory() + "/" + source + Utils::GLShaderStageOpenGLFileExtension(stage);
				CC_CORE_INFO("m_FilePath = %s\n", m_FilePath.c_str());

				std::string source_text = FileLog::ReadFile(m_FilePath);
				CC_ASSERT(!source_text.empty(), "ASSERT: CompileOrGetOpenGLBinaries %s is null\n", m_FilePath.c_str());

				shaderc::SpvCompilationResult module = compiler.CompileGlslToSpv(source_text, Utils::GLShaderStageToShaderC(stage), m_FilePath.c_str(), options);
				if (module.GetCompilationStatus() != shaderc_compilation_status_success)
				{
					CC_CORE_ERROR("%s\n", module.GetErrorMessage().c_str());
					CC_ASSERT(false, "ASSERT: module.GetCompilationStatus() error.\n");
				}

				shaderData[stage] = std::vector<uint32_t>(module.cbegin(), module.cend());

				std::ofstream out(cachedPath, std::ios::out | std::ios::binary);
				if (out.is_open())
				{
					auto &data = shaderData[stage];
					out.write((char *)data.data(), data.size() * sizeof(uint32_t));
					out.flush();
					out.close();
				}
			}
		}

	}

	void OpenGLShader::CreateProgram()
	{
		GLuint program = glCreateProgram();

		std::vector<GLuint> shaderIDs;
		std::unordered_map<GLenum, std::vector<uint32_t>>::iterator it = m_OpenGLSPIRV.begin();
		for(it; it != m_OpenGLSPIRV.end(); it++)
		{
			GLenum stage = it->first;
			std::vector<uint32_t> &spirv = it->second;

            GLuint shaderID = glCreateShader(stage);
            shaderIDs.emplace_back(shaderID);
			glShaderBinary(1, &shaderID, GL_SHADER_BINARY_FORMAT_SPIR_V, spirv.data(), spirv.size() * sizeof(uint32_t));
			glSpecializeShader(shaderID, "main", 0, nullptr, nullptr);
			glAttachShader(program, shaderID);
		}

		glLinkProgram(program);

		CheckCompileErrors(program, "PROGRAM");

        for (auto id : shaderIDs)
        {
            glDetachShader(program, id);
            glDeleteShader(id);
        }
		m_RendererID = program;
	}

	void OpenGLShader::CompileOnly(const std::string &vertexFile, const std::string &fragmentFile){
		std::string vertexCode;
		std::string fragmentCode;
		std::string geometryCode;

		vertexCode = FileLog::ReadFile(vertexFile);
		fragmentCode = FileLog::ReadFile(fragmentFile);

		const char *vShaderCode = vertexCode.c_str();
		const char *fShaderCode = fragmentCode.c_str();
		
		unsigned int vertex, fragment;
		// vertex shader
		vertex = glCreateShader(GL_VERTEX_SHADER);
		glShaderSource(vertex, 1, &vShaderCode, NULL);
		glCompileShader(vertex);
		CheckCompileErrors(vertex, "VERTEX");

		// fragment Shader
		fragment = glCreateShader(GL_FRAGMENT_SHADER);
		glShaderSource(fragment, 1, &fShaderCode, NULL);
		glCompileShader(fragment);
		CheckCompileErrors(fragment, "FRAGMENT");
		// if geometry shader is given, compile geometry shader
		// unsigned int geometry;
		// shader Program
		m_RendererID = glCreateProgram();
		// printf("shader Program ID=%d\n", ID);
		glAttachShader(m_RendererID, vertex);
		glAttachShader(m_RendererID, fragment);

		glLinkProgram(m_RendererID);

		CheckCompileErrors(m_RendererID, "PROGRAM");
		// delete the shaders as they're linked into our program now and no longer necessery
		glDeleteShader(vertex);
		glDeleteShader(fragment);
	}

	void OpenGLShader::CheckCompileErrors(unsigned int shader, std::string type)
	{
		GLint success;
		GLchar infoLog[1024];
		if (type != "PROGRAM")
		{
			glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
			if (!success)
			{
				glGetShaderInfoLog(shader, 1024, NULL, infoLog);
				std::cout << "ERROR::SHADER_COMPILATION_ERROR of type: " << type << "\n"
						  << infoLog << "\n -- --------------------------------------------------- -- " << std::endl;
				CC_ASSERT(false, "ASSERT: CheckCompileErrors %s\n", type.c_str());
			}
		}
		else
		{
			glGetProgramiv(shader, GL_LINK_STATUS, &success);
			if (!success)
			{
				glGetProgramInfoLog(shader, 1024, NULL, infoLog);
				std::cout << "ERROR::PROGRAM_LINKING_ERROR of type: " << type << "\n"
						  << infoLog << "\n -- --------------------------------------------------- -- " << std::endl;
				glDeleteProgram(shader);
				CC_ASSERT(false, "ASSERT: CheckCompileErrors %s\n", type.c_str());
			}
		}
	}

	void OpenGLShader::Bind() const
	{
		glUseProgram(m_RendererID);
	}

	void OpenGLShader::Unbind() const
	{
		glUseProgram(0);
	}

	void OpenGLShader::SetInt(const std::string &name, int value)
	{
		UploadUniformInt(name, value);
	}

	void OpenGLShader::SetIntArray(const std::string &name, int *values, uint32_t count)
	{
		UploadUniformIntArray(name, values, count);
	}

	void OpenGLShader::SetFloat(const std::string &name, float value)
	{

		UploadUniformFloat(name, value);
	}

	void OpenGLShader::SetFloat2(const std::string &name, const glm::vec2 &value)
	{

		UploadUniformFloat2(name, value);
	}

	void OpenGLShader::SetFloat3(const std::string &name, const glm::vec3 &value)
	{

		UploadUniformFloat3(name, value);
	}

	void OpenGLShader::SetFloat4(const std::string &name, const glm::vec4 &value)
	{

		UploadUniformFloat4(name, value);
	}

	void OpenGLShader::SetMat4(const std::string &name, const glm::mat4 &value)
	{

		UploadUniformMat4(name, value);
	}

	void OpenGLShader::UploadUniformInt(const std::string &name, int value)
	{
		GLint location = glGetUniformLocation(m_RendererID, name.c_str());
		glUniform1i(location, value);
	}

	void OpenGLShader::UploadUniformIntArray(const std::string &name, int *values, uint32_t count)
	{
		GLint location = glGetUniformLocation(m_RendererID, name.c_str());
		glUniform1iv(location, count, values);
	}

	void OpenGLShader::UploadUniformFloat(const std::string &name, float value)
	{
		GLint location = glGetUniformLocation(m_RendererID, name.c_str());
		glUniform1f(location, value);
	}

	void OpenGLShader::UploadUniformFloat2(const std::string &name, const glm::vec2 &value)
	{
		GLint location = glGetUniformLocation(m_RendererID, name.c_str());
		glUniform2f(location, value.x, value.y);
	}

	void OpenGLShader::UploadUniformFloat3(const std::string &name, const glm::vec3 &value)
	{
		GLint location = glGetUniformLocation(m_RendererID, name.c_str());
		glUniform3f(location, value.x, value.y, value.z);
	}

	void OpenGLShader::UploadUniformFloat4(const std::string &name, const glm::vec4 &value)
	{
		GLint location = glGetUniformLocation(m_RendererID, name.c_str());
		glUniform4f(location, value.x, value.y, value.z, value.w);
	}

	void OpenGLShader::UploadUniformMat3(const std::string &name, const glm::mat3 &matrix)
	{
		GLint location = glGetUniformLocation(m_RendererID, name.c_str());
		glUniformMatrix3fv(location, 1, GL_FALSE, glm::value_ptr(matrix));
	}

	void OpenGLShader::UploadUniformMat4(const std::string &name, const glm::mat4 &matrix)
	{
		GLint location = glGetUniformLocation(m_RendererID, name.c_str());
		glUniformMatrix4fv(location, 1, GL_FALSE, glm::value_ptr(matrix));
	}

	uint32_t OpenGLShader::GetProgramID() const
	{
		return m_RendererID;
	}
}
